import time
import control
import math
import scipy.io
from numpy.linalg import inv
import scipy.integrate
from scipy.spatial.transform import Rotation
import csv
import numpy.linalg as LA
import numpy as np
import cvxpy as cp
import matplotlib.pyplot as plt





# %%
class Terminal_set:
    def __init__(self, Hx, Hu, K, Ak, h):
        # 设定时间步长
        self.dt = 0.01
        # 初始化参数：状态/输入约束矩阵、反馈增益、闭环系统矩阵、约束向量
        self.Hx = Hx
        self.Hu = Hu
        self.K = K
        self.Ak = Ak
        # 组合状态和输入约束矩阵
        self.H = np.block([[Hu, np.zeros((Hu.shape[0], Hx.shape[1]))],[np.zeros((Hx.shape[0], Hu.shape[1])), Hx]])
        self.Nc = self.H.shape[0]
        self.Nx = Ak.shape[1]
        self.h = h
        self.K_aug = np.vstack((-K, np.eye(self.Nx)))
        self.maxiter = 200
        # 计算终端约束集
        self.Xf = self.terminal_set_cal()
        # 去除冗余约束
        self.Xf_nr = self.remove_redundancy()
        # 验证输入约束
        self.test_input_inbound(0.15)
    # 通过迭代线性规划寻找满足约束的最大不变集。
    def terminal_set_cal(self):
        Ainf = np.zeros([0, self.Nx])
        binf = np.zeros([0, 1])
        Ft = np.eye(self.Nx)
        self.C = self.H@self.K_aug

        for t in range(self.maxiter):
            Ainf = np.vstack((Ainf, self.C@Ft))
            binf = np.vstack((binf, self.h))

            Ft = self.Ak@Ft
            fobj = self.C@Ft
            violation = False
            for i in range(self.Nc):
                val, x = self.solve_linprog(fobj[i, :], Ainf, binf)
                if val > self.h[i]:
                    violation = True
                    break
            if not violation:
                return [Ainf, binf]

    def solve_linprog(self, obj, Ainf, binf):
        x = cp.Variable((self.Nx, 1))
        objective = cp.Maximize(obj@x)
        constraints = [Ainf@x <= binf]
        linear_program = cp.Problem(objective, constraints)

        result = linear_program.solve(verbose=False)
        return result, x.value
    # 去除冗余约束，简化终端集。
    def remove_redundancy(self):
        A_inf, Binf = self.Xf
        Ainf_nr, binf_nr = A_inf.copy(), Binf.copy()
        i = 0
        while i < Ainf_nr.shape[0]:
            obj = Ainf_nr[i, :]
            binf_temp = binf_nr.copy()
            binf_temp[i] += 1
            val, x = self.solve_linprog(obj, Ainf_nr, binf_temp)
            if val < binf_nr[i] or val == binf_nr[i]:
                Ainf_nr = np.delete(Ainf_nr, i, 0)
                binf_nr = np.delete(binf_nr, i, 0)
            else:
                i += 1
        return [Ainf_nr, binf_nr]
    # 验证终端集内控制输入是否满足约束。
    def test_input_inbound(self,u_limit):
        A_inf,b_inf=self.Xf_nr
        violation =False
        for i in range(4):
            x = cp.Variable((12, 1))
            u = cp.Variable((4, 1))
            cost=0
            constr = []
            constr.append(A_inf@x[:,0] <= b_inf.squeeze())
            constr.append(u[:, 0]==-self.K@x[:,0])
            cost=u[i, 0]
            problem = cp.Problem(cp.Maximize(cost), constr)
            problem.solve()
            print('Input u',i,'<',problem.value)
            if problem.value >u_limit:
                violation =True
        if violation ==False:
            print('Input inbound')





# %%
class Trajectory:
    """
    轨迹生成器类
    用于生成不同类型的参考轨迹，包括菱形、直线、悬停和圆形轨迹
    """
    def __init__(self, input_traj="hover"):

        self.heading = np.zeros((2))
        self.yaw = 0
        if input_traj == "diamond":
            self.input_traj = self.diamond
        elif input_traj == "oneline":
            self.input_traj = self.oneline
        elif input_traj == "hover":
            self.input_traj = self.hover
        elif input_traj == "circle":
            self.input_traj = self.circle

    def circle(self, t):
        T = 14
        radius = 6
        dt = 0.01
        if t + dt > T:
            pos = np.array([0, 0, 2.5]).reshape([-1, 1])
            vel = np.array([0, 0, 0]).reshape([-1, 1])
            acc = np.array([0, 0, 0]).reshape([-1, 1])
        else:
            angle, _, _ = self.tj_from_line(0, 2 * np.pi, T, t)
            angle2, _, _ = self.tj_from_line(0, 2 * np.pi, T, t + dt)
            angle3, _, _ = self.tj_from_line(0, 2 * np.pi, T, t + 2 * dt)
            pos = np.array([radius * (np.cos(angle) - 1), radius * np.sin(angle), 2.5 * angle / (2 * np.pi)])
            pos2 = np.array([radius * (np.cos(angle2) - 1), radius * np.sin(angle2), 2.5 * angle2 / (2 * np.pi)])
            pos3 = np.array([radius * (np.cos(angle3) - 1), radius * np.sin(angle3), 2.5 * angle3 / (2 * np.pi)])
            vel = (pos2 - pos) / dt
            vel2 = (pos3 - pos2) / dt
            acc = (vel2 - vel) / dt
        pos = pos.reshape([-1, 1])
        vel = vel.reshape([-1, 1])
        acc = acc.reshape([-1, 1])
        return pos, vel, acc

    def diamond(self, t):
        T1, T2, T3, T4 = 3.5, 3.5, 3.5, 3.5
        points = []
        points.append(np.zeros((3, 1)))


        points.append(np.array([[0], [np.sqrt(2)], [np.sqrt(2)]]))
        points.append(np.array([[1], [0], [2 * np.sqrt(2)]]))
        points.append(np.array([[1], [-np.sqrt(2)], [np.sqrt(2)]]))
        points.append(np.array([[1], [0], [0]]))

        if (0 < t) and (t <= T1):
            pos, vel, acc = self.tj_from_line(points[0], points[1], T1, t)
        elif (T1 < t) and (t <= (T1 + T2)):
            pos, vel, acc = self.tj_from_line(points[1], points[2], T2, t - T1)
        elif ((T1 + T2) < t) and (t <= (T1 + T2 + T3)):
            pos, vel, acc = self.tj_from_line(
                points[2], points[3], T3, t - (T1 + T2))
        elif ((T1 + T2 + T3) < t) and (t <= (T1 + T2 + T3 + T4)):
            pos, vel, acc = self.tj_from_line(
                points[3], points[4], T4, t - (T1 + T2 + T3))
        elif (t > (T1 + T2 + T3 + T4)):
            pos, vel, acc = points[4], np.zeros((3, 1)), np.zeros((3, 1))
        else:
            pos, vel, acc = np.zeros((3, 1)), np.zeros((3, 1)), np.zeros((3, 1))
        return pos, vel, acc

    def oneline(self, t):
        T1 = 5.0
        points = []
        points.append(np.zeros((3, 1)))
        points.append(np.array([[0], [1], [1]]))
        if (0 < t) and (t <= T1):
            pos, vel, acc = self.tj_from_line(points[0], points[1], T1, t)
        else:
            pos, vel, acc = points[-1], np.zeros((3, 1)), np.zeros((3, 1))
        return pos, vel, acc

    def hover(self, t):
        pos, vel, acc = np.zeros((3, 1)), np.zeros((3, 1)), np.zeros((3, 1))
        return pos, vel, acc

    def tj_from_line(self, start_pos, end_pos, time_ttl, t_c):
        v_max = (end_pos - start_pos) * 2 / time_ttl
        if (t_c >= 0 and t_c < time_ttl / 2):
            vel = v_max * t_c / (time_ttl / 2)
            pos = start_pos + t_c * vel / 2
            acc = v_max / (time_ttl / 2)
        elif (t_c >= time_ttl / 2 and t_c <= time_ttl):
            vel = v_max * (time_ttl - t_c) / (time_ttl / 2)
            pos = end_pos - (time_ttl - t_c) * vel / 2
            acc = - v_max / (time_ttl / 2)
        else:
            if (type(start_pos) == int) or (type(start_pos) == float):
                pos, vel, acc = 0, 0, 0
            else:
                pos, vel, acc = np.zeros(start_pos.shape), np.zeros(
                    start_pos.shape), np.zeros(start_pos.shape)
        return pos, vel, acc

    def get_yaw(self, vel):
        if np.allclose(vel, np.zeros((3, 1))):
            vel += 1e-5
        curr_heading = vel / LA.norm(vel)
        prev_heading = self.heading
        cosine = max(-1, min(np.dot(prev_heading, curr_heading), 1))
        dyaw = np.arccos(cosine)
        norm_v = np.cross(prev_heading, curr_heading)
        self.yaw += np.sign(norm_v) * dyaw

        if self.yaw > np.pi:
            self.yaw -= 2 * np.pi
        if self.yaw < -np.pi:
            self.yaw += 2 * np.pi
        self.heading = curr_heading
        yaw_dot = max(-30, min(dyaw / 0.005, 30))
        return self.yaw, yaw_dot

    def get_des_state(self, t):
        pos, vel, acc = self.input_traj(t)
        # yaw, yaw_dot = self.get_yaw(vel[:2].flatten())
        des_state = {'x': pos, 'v': vel,
                     'x_ddt': acc, 'yaw': 0, 'yaw_dot': 0}
        return des_state


# %%
class Visualizer:
    def __init__(self, simu_time, simu_freq, ctrl_freq, real_trajectory, des_trajectory) -> None:
        self.simu_time = simu_time
        self.simu_freq = simu_freq
        self.ctrl_freq = ctrl_freq
        self.real_trajectory = real_trajectory
        self.des_trajectory = des_trajectory
        self.real_trajectory['x'] = np.array(self.real_trajectory['x'])
        self.real_trajectory['y'] = np.array(self.real_trajectory['y'])
        self.real_trajectory['z'] = np.array(self.real_trajectory['z'])
        self.des_trajectory['x'] = np.array(self.des_trajectory['x'])
        self.des_trajectory['y'] = np.array(self.des_trajectory['y'])
        self.des_trajectory['z'] = np.array(self.des_trajectory['z'])
        self.file_data = open('./data.csv', 'a')

    def plot_3d_trajectory(self):
        """Plot static 3D trajectory"""
        # 绘制3D轨迹（真实 vs 期望）
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')

        # Plot real trajectory
        ax.plot(self.real_trajectory['x'],
                self.real_trajectory['y'],
                self.real_trajectory['z'],
                'r-',
                label='Real Trajectory')

        # Plot desired trajectory
        ax.plot(self.des_trajectory['x'],
                self.des_trajectory['y'],
                self.des_trajectory['z'],
                'b--',
                label='Desired Trajectory')

        # Mark start and end points
        ax.scatter(self.real_trajectory['x'][0],
                   self.real_trajectory['y'][0],
                   self.real_trajectory['z'][0],
                   c='g',
                   marker='o',
                   label='Start')
        ax.scatter(self.real_trajectory['x'][-1],
                   self.real_trajectory['y'][-1],
                   self.real_trajectory['z'][-1],
                   c='r',
                   marker='s',
                   label='End')

        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('Z')
        ax.set_title('Quadrotor 3D Trajectory')
        ax.legend()

        # Set axis limits
        ax.set_xlim([-6, 1])
        ax.set_ylim([-3, 3])
        ax.set_zlim([0, 2.0])

        # Set the viewing angle
        ax.view_init(35, 35)

        plt.show()

    # 分坐标轴绘制位置跟踪性能
    def plot_tracking_performance(self):
        """Plot tracking performance in separate coordinates"""
        t = np.arange(0, self.simu_time, 1 / self.simu_freq)

        plt.figure(figsize=(10, 6))
        plt.subplot(311)
        plt.plot(t, self.real_trajectory['x'], 'r-', label='x_real')
        plt.plot(t, self.des_trajectory['x'], 'b--', label='x_des')
        plt.ylabel('X Position')
        plt.legend()
        plt.grid(True)

        plt.subplot(312)
        plt.plot(t, self.real_trajectory['y'], 'r-', label='y_real')
        plt.plot(t, self.des_trajectory['y'], 'b--', label='y_des')
        plt.ylabel('Y Position')
        plt.legend()
        plt.grid(True)

        plt.subplot(313)
        plt.plot(t, self.real_trajectory['z'], 'r-', label='z_real')
        plt.plot(t, self.des_trajectory['z'], 'b--', label='z_des')
        plt.xlabel('Time (s)')
        plt.ylabel('Z Position')
        plt.legend()
        plt.grid(True)

        plt.tight_layout()
        plt.show()

    def record_tracking_data(self):
        writer = csv.writer(self.file_data)
        writer.writerow(self.real_trajectory['x'])
        writer.writerow(self.des_trajectory['x'])
# %% %

def quat_dot(quat, omega):
    """
    Parameters:
        quat, [i,j,k,w]
        omega, angular velocity of body in body axes

    Returns
        quat_dot, [i,j,k,w]

    """
    # Adapted from "Quaternions And Dynamics" by Basile Graf.
    (q0, q1, q2, q3) = (quat[0], quat[1], quat[2], quat[3])
    G = np.array([[q3, q2, -q1, -q0],
                  [-q2, q3, q0, -q1],
                  [q1, -q0, q3, -q2]])
    quat_dot = 0.5 * G.T @ omega
    # Augment to maintain unit quaternion.
    quat_err = np.sum(quat ** 2) - 1
    quat_err_grad = 2 * quat
    quat_dot = quat_dot - quat_err * quat_err_grad
    return quat_dot


class Quadrotor():
    """
    Quadrotor forward dynamics model.
    """

    def __init__(self):
        # 物理参数：质量、惯性矩、旋翼参数等
        self.mass = 0.030  # kg
        self.Ixx = 1.43e-5  # kg*m^2
        self.Iyy = 1.43e-5  # kg*m^2
        self.Izz = 2.89e-5  # kg*m^2
        self.arm_length = 0.046  # meters
        self.rotor_speed_min = 0  # rad/s
        self.rotor_speed_max = 2500  # rad/s
        self.k_thrust = 2.3e-08  # N/(rad/s)**2
        self.k_drag = 7.8e-11  # Nm/(rad/s)**2

        # Additional constants.
        self.inertia = np.diag(
            np.array([self.Ixx, self.Iyy, self.Izz]))  # kg*m^2
        self.g = 9.81  # m/s^2

        # Precomputes
        k = self.k_drag / self.k_thrust
        L = self.arm_length
        self.to_TM = np.array([[1, 1, 1, 1],
                               [0, L, 0, -L],
                               [-L, 0, L, 0],
                               [k, -k, k, -k]])
        self.inv_inertia = inv(self.inertia)
        self.weight = np.array([0, 0, -self.mass * self.g])
        self.t_step = 0.01

        # Initialize state
        self.state = _unpack_state(np.zeros(13))

    def reset(self, position=[0, 0, 0], yaw=0, pitch=0, roll=0):
        '''
        state is a 13 dimensional vector
            postion*3 velocity*3 attitude(quaternion)*4 angular velocity*3
        state = [x y z dx dy dz qw qx qy qz r p q]
        dot_state = [dx dy dz ddx ddy ddz dqw dqx dqy dqz dr dp dq]
        '''
        s = np.zeros(13)
        s[0] = position[0]
        s[1] = position[1]
        s[2] = position[2]
        r = Rotation.from_euler('zxy', [yaw, roll, pitch], degrees=True)
        quat = r.as_quat()
        s[6] = quat[0]
        s[7] = quat[1]
        s[8] = quat[2]
        s[9] = quat[3]
        # the unassigned values of s are zeros
        self.state = _unpack_state(s)
        return self.state

    # 应用控制输入，通过Runge-Kutta积分更新状态
    def step(self, cmd_rotor_speeds):
        '''
        Considering the max and min of rotor speeds
        action is a 4 dimensional vector: conmmand rotor speeds
        action = [w1, w2, w3, w4]
        '''
        rotor_speeds = np.clip(
            cmd_rotor_speeds, self.rotor_speed_min, self.rotor_speed_max)
        rotor_thrusts = self.k_thrust * rotor_speeds ** 2

        '''
        Next, [w1, w2, w3, w4] into [F Mx My Mz]
        '''
        TM = self.to_TM @ rotor_thrusts
        T = TM[0]  # u1
        M = TM[1:]  # u2


        def s_dot_fn(t, s):
            return self._s_dot_fn(t, s, T, M)

        '''
        The next state can be obtained through integration （Runge-Kutta）
        '''
        s = _pack_state(self.state)
        sol = scipy.integrate.solve_ivp(s_dot_fn, (0, self.t_step), s, first_step=self.t_step)
        s = sol['y'][:, -1]
        # turn state back to dict
        self.state = _unpack_state(s)

        # Re-normalize unit quaternion.
        reward = 0
        done = 0
        info = {}
        return self.state, reward, done, info

    def _s_dot_fn(self, t, s, u1, u2):
        """
        Compute derivative of state for quadrotor given fixed control inputs as
        an autonomous ODE.
        """

        state = _unpack_state(s)
        # page 73
        # Position derivative.
        x_dot = state['v']

        # Velocity derivative.
        F = u1 * Quadrotor.rotate_k(state['q'])
        v_dot = (self.weight + F) / self.mass

        # Orientation derivative.
        q_dot = quat_dot(state['q'], state['w'])

        # Angular velocity derivative. page 26 Equation 4
        omega = state['w']
        omega_hat = Quadrotor.hat_map(omega)
        w_dot = self.inv_inertia @ (u2 - omega_hat @ (self.inertia @ omega))

        # Pack into vector of derivatives.
        s_dot = np.zeros((13,))
        s_dot[0:3] = x_dot
        s_dot[3:6] = v_dot
        s_dot[6:10] = q_dot
        s_dot[10:13] = w_dot

        return s_dot

    # 获取离散化后的线性时不变系统矩阵（A, B）
    def get_dLTI(self, dt):
        num_x = 12
        Ac = np.array([[0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
                       [0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
                       [0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
                       [0, 0, 0, 0, 0, 0, 0, self.g, 0, 0, 0, 0],
                       [0, 0, 0, 0, 0, 0, (-self.g), 0, 0, 0, 0, 0],
                       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                       [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
                       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
                       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
                       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], ])
        Bc = np.array([[0, 0, 0, 0],
                       [0, 0, 0, 0],
                       [0, 0, 0, 0],
                       [0, 0, 0, 0],
                       [0, 0, 0, 0],
                       [1 / self.mass, 0, 0, 0],
                       [0, 0, 0, 0],
                       [0, 0, 0, 0],
                       [0, 0, 0, 0],
                       [0, 1 / self.Ixx, 0, 0],
                       [0, 0, 1 / self.Iyy, 0],
                       [0, 0, 0, 1 / self.Izz]])
        Cc = np.eye(num_x)
        Dc = np.zeros((num_x, 4))
        sysc = control.ss(Ac, Bc, Cc, Dc)
        # Discretization
        sysd = control.sample_system(sysc, dt, method='bilinear')
        return sysd.A, sysd.B

    @classmethod
    def rotate_k(cls, q):
        """
        Rotate the unit vector k by quaternion q. This is the third column of
        the rotation matrix associated with a rotation by q.
        """
        return np.array([2 * (q[0] * q[2] + q[1] * q[3]),
                         2 * (q[1] * q[2] - q[0] * q[3]),
                         1 - 2 * (q[0] ** 2 + q[1] ** 2)])

    @classmethod
    def hat_map(cls, s):
        """
        Given vector s in R^3, return associate skew symmetric matrix S in R^3x3
        """
        return np.array([[0, -s[2], s[1]],
                         [s[2], 0, -s[0]],
                         [-s[1], s[0], 0]])


def _pack_state(state):
    """
    Convert a state dict to Quadrotor's private internal vector representation.
    """
    s = np.zeros((13,))
    s[0:3] = state['x'].squeeze()
    s[3:6] = state['v'].squeeze()
    s[6:10] = state['q'].squeeze()
    s[10:13] = state['w'].squeeze()
    return s


def _unpack_state(s):
    """
    Convert Quadrotor's private internal vector representation to a state dict.
    """
    state = {'x': s[0:3], 'v': s[3:6], 'q': s[6:10], 'w': s[10:13]}
    return state


# %%
class Controller:
    def __init__(self, traj, ctrl_freq):
        self.quad_model = Quadrotor()
        self.traj = traj
        self.ctrl_freq = ctrl_freq
        self.mass = self.quad_model.mass
        self.g = self.quad_model.g
        self.arm_length = self.quad_model.arm_length
        self.rotor_speed_min = 0
        self.rotor_speed_max = 2500
        self.k_thrust = self.quad_model.k_thrust
        self.k_drag = self.quad_model.k_drag
        self.to_TM = self.quad_model.to_TM
        self.inertia = np.diag(
            [self.quad_model.Ixx, self.quad_model.Iyy, self.quad_model.Izz])

    def generate_control_input(self, u):
        cmd_motor_speeds = np.zeros((4,))
        cmd_thrust = 0
        cmd_moment = np.zeros((3,))
        # cmd_q = np.zeros((4,))

        F = (np.linalg.inv(self.to_TM) @ u).astype(float)

        for i in range(4):
            if F[i] < 0:
                F[i] = 0
                cmd_motor_speeds[i] = self.rotor_speed_min
            cmd_motor_speeds[i] = np.sqrt(F[i] / self.k_thrust)
            if cmd_motor_speeds[i] > self.rotor_speed_max:
                cmd_motor_speeds[i] = self.rotor_speed_max

        cmd_thrust = u[0]  # thrust
        cmd_moment[0] = u[1]  # moment about p
        cmd_moment[1] = u[2]  # moment about q
        cmd_moment[2] = u[3]  # moment about r

        control_input = {'cmd_motor_speeds': cmd_motor_speeds,
                         'cmd_thrust': cmd_thrust,
                         'cmd_moment': cmd_moment}
        return control_input

    def state2x(self, state):
        x = state.get('x').flatten()
        v = state.get('v').flatten()
        try:
            q = state.get('q').flatten()
            w = state.get('w').flatten()
            euler_ang = self.euler_from_quaternion(q[0], q[1], q[2], q[3])
        except:
            euler_ang = np.zeros(3)
            euler_ang[2] = state.get('yaw')
            w = np.zeros(3)
            w[2] = state.get('yaw_dot')

        x_init = np.block([x, v, euler_ang, w])
        return x_init

    def euler_from_quaternion(self, x, y, z, w):
        """
        Convert a quaternion into euler angles (roll, pitch, yaw)
        roll is rotation around x in radians (counterclockwise)
        pitch is rotation around y in radians (counterclockwise)
        yaw is rotation around z in radians (counterclockwise)
        """
        t0 = +2.0 * (w * x + y * z)
        t1 = +1.0 - 2.0 * (x * x + y * y)
        roll_x = math.atan2(t0, t1)

        t2 = +2.0 * (w * y - z * x)
        t2 = +1.0 if t2 > +1.0 else t2
        t2 = -1.0 if t2 < -1.0 else t2
        pitch_y = math.asin(t2)

        t3 = +2.0 * (w * z + x * y)
        t4 = +1.0 - 2.0 * (y * y + z * z)
        yaw_z = math.atan2(t3, t4)

        return np.array([roll_x, pitch_y, yaw_z])  # in radians



class Linear_MPC(Controller):
    # traj 表示参考轨迹，ctrl_freq 是控制频率，use_obsv 是一个布尔值，用于决定是否使用观测器。
    def __init__(self, traj, ctrl_freq, use_obsv=False):
        Controller.__init__(self, traj, ctrl_freq)
        self.ctrl_freq = ctrl_freq
        # self.dt 是控制周期，即控制频率的倒数。
        self.dt = 1 / self.ctrl_freq

        # 获取离散化的系统状态空间矩阵
        self.Ad, self.Bd = self.quad_model.get_dLTI(self.dt)

        # 初始化重力补偿向量
        self.gravity = np.zeros([12, ])
        self.gravity[5] = self.dt * self.g  # discritized
        self.use_obsv = use_obsv

        # 定义干扰输入矩阵
        self.B_dist = np.block([[np.zeros((3, 3))],
                                [np.eye(3)],
                                [np.zeros((3, 3))],
                                [np.zeros((3, 3))]])

        # 初始化状态观测器和速度观测器
        self.disturbance_observer = Luenberger_Observer(self.Ad, self.Bd, self.B_dist, C=np.eye(12),
                                                        C_dist=np.zeros((12, 3)), gravity=self.gravity, load_L=False)

        # 定义C矩阵用于观测位置和姿态
        C_obs = np.zeros((6, 12))
        C_obs[:3, :3] = np.eye(3)
        C_obs[3:, 6:9] = np.eye(3)
        self.vel_observer = Luenberger_Observer(self.Ad, self.Bd, self.B_dist, C=C_obs, C_dist=np.zeros((6, 3)),
                                                gravity=self.gravity, load_L=True)

        # 存储实际状态和观测状态的列表
        self.x_real = [[], [], []]
        self.x_obsv = [[], [], []]
        self.d_hat_list = []


        # MPC控制器参数设置，self.N 是预测时域的长度。
        self.N = 5

        # 状态权重矩阵Q：调整位置、速度、姿态的跟踪精度
        self.Q = np.diag([80, 80, 100, 80, 80, 100, 50, 50, 50, 50, 50, 50])

        # 控制输入权重矩阵R：平衡控制能量消耗
        self.R = np.diag([50, 80, 80, 80])

        # 计算终端代价矩阵P和反馈增益K
        self.P, self.K = self.get_terminal_cost(self.Ad, self.Bd, self.Q, self.R)

        # 闭环系统矩阵
        self.Ak = self.Ad - self.Bd @ self.K

        # 定义状态约束矩阵Hx（限制姿态角和速度）
        self.Hx = np.array([
            [0, 0, 0, 0, 0, 0, 0, 1.0, 0, 0, 0, 0],  # roll pitch约束
            [0, 0, 0, 0, 0, 0, -1.0, 0, 0, 0, 0, 0],
            [0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0], # 速度约束
            [0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, -1.0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 1.0, 0, 0, 0, 0, 0],
            [0, 0, 0, -1.0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, -1.0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, -1.0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 1.0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, -1.0, 0, 0, 0],
        ])

        # 定义输入约束矩阵（限制推力和力矩）
        self.Hu = np.array([[1 / self.mass, 0, 0, 0],
                            [-1 / self.mass, 0, 0, 0]
                            ])  # z acc constraints
        self.Hu1 = np.array([[1 / self.mass, 0, 0, 0],
                             [-1 / self.mass, 0, 0, 0],
                             [0, 1.0, 0, 0],
                             [0, -1.0, 0, 0],
                             [0, 0, 1.0, 0],
                             [0, 0, -1.0, 0],
                             [0, 0, 0, 1.0],
                             [0, 0, 0, -1.0]
                             ])
        self.h = np.array([[0.5 * self.g],  # z acc constraints
                           [0.5 * self.g],
                           [0.5],
                           [0.5],
                           [2.0],  # velocity constraints
                           [2.0],
                           [2.0],
                           [0.5],
                           [0.5],
                           [2.0],
                           [2.0],
                           [2.0],
                           [0.5],
                           [0.5]
                           ])
        self.h1 = np.array([[1.5 * self.g],  # z acc constraints
                            [-0.5 * self.g],
                            [0.15],
                            [0.15],
                            [0.15],
                            [0.15],
                            [0.15],
                            [0.15],
                            [0.5],
                            [0.5],
                            [2.0],  # velocity constraints
                            [2.0],
                            [2.0],
                            [0.5],
                            [0.5],
                            [2.0],
                            [2.0],
                            [2.0],
                            [0.5],
                            [0.5]
                            ])
        self.terminal_set = Terminal_set(
            self.Hx, self.Hu, self.K, self.Ak, self.h)

        self.Xf_nr = self.terminal_set.Xf_nr

        self.terminal_set.test_input_inbound(u_limit=0.15)
        # raise NotImplementedError()
        self.x_real = [[], [], []]
        self.x_obsv = [[], [], []]

    def control(self, cur_time, obs_state):
        # 获取期望状态和计算位置误差
        des_state = self.traj.get_des_state(cur_time)
        error_pos = (des_state.get('x') - obs_state.get('x').reshape(3, 1)).flatten()


        # 获取系统当前状态
        x_sys = self.state2x(obs_state)
        x_obs_dist = self.disturbance_observer.x_hat.flatten()
        x_obs_vel = self.vel_observer.x_hat.flatten()
        y = np.block([x_sys[:3], x_sys[6:9]])

        x_init = x_sys
        d_hat = self.disturbance_observer.d_hat
        self.d_hat_list.append(d_hat[1])
        for i in range(3, 6):
            self.x_real[i - 3].append(x_sys[i])
            self.x_obsv[i - 3].append(x_obs_vel[i])

        # 定义MPC优化变量
        x = cp.Variable((12, self.N + 1))  # 状态变量
        u = cp.Variable((4, self.N))       # 控制输入变量
        cost = 0
        constr = []

        # 在预测时域内构建代价函数和约束
        mpc_time = cur_time
        desired_x = []
        for k in range(self.N + 1):
            mpc_time += k * self.dt

            # 获取未来参考轨迹
            des_state_ahead = self.traj.get_des_state(mpc_time)
            x_ref_k = self.state2x(des_state_ahead)
            desired_x.append(x_ref_k)
            if k == self.N:
                # 添加终端代价
                cost += cp.quad_form(x[:, self.N] - x_ref_k, self.P)

                # 终端约束
                if self.use_obsv == False:
                    constr.append(
                        self.Xf_nr[0] @ (x[:, self.N] - x_ref_k) <= self.Xf_nr[1].squeeze())
                break

            # 状态跟踪代价
            cost += cp.quad_form(x[:, k] - x_ref_k, self.Q)

            # 控制输入代价
            u_ref_k = np.array([self.mass * self.g, 0, 0, 0])
            cost += cp.quad_form(u[:, k] - u_ref_k, self.R)

            # 添加状态和输入约束
            constr.append(self.Hx @ x[:, k] <= self.h1[self.Hu1.shape[0]:].squeeze())
            constr.append(self.Hu1 @ u[:, k] <= self.h1[:self.Hu1.shape[0]].squeeze())
            gravity = np.zeros([12, ])
            gravity[5] = self.dt * self.g  # discritized

            # 系统动力学约束
            constr.append(x[:, k + 1] == self.Ad @ x[:, k] +
                          self.Bd @ u[:, k] - self.gravity + (self.B_dist @ d_hat).flatten())

        # 初始状态约束
        constr.append(x[:, 0] == x_init)

        # 求解优化问题
        problem = cp.Problem(cp.Minimize(cost), constr)
        problem.solve(verbose=False)

        # 获取最优控制输入
        u = u[:, 0].value
        print(u)

        # 更新观测器
        if self.use_obsv:
            self.disturbance_observer.update(u, x_sys)
            self.vel_observer.update(u, y)

        # 生成实际的电机控制指令
        control_input = self.generate_control_input(u)
        return control_input, error_pos

    def get_terminal_cost(self, Ad, Bd, Q, R):
        P, L, G = control.dare(Ad, Bd, Q, R)
        return P, G


class Luenberger_Observer:
    def __init__(self, A, B, B_dist, C, C_dist, gravity, load_L=False):
        self.x_hat = np.zeros((12, 1)) + 0.1
        self.d_hat = np.zeros((3, 1))
        self.A_sq = np.block([[A, B_dist],
                              [np.zeros((3, 12)), np.eye(3)]])
        self.B_sq = np.block([[B],
                              [np.zeros((3, 4))]])
        self.C_sq = np.block([C, C_dist])
        self.gravity = np.block([[gravity.reshape(-1, 1)],
                                 [np.zeros((3, 1))]])

        # 计算观测器增益矩阵L
        self.L = self.compute_observer_gain()

    def compute_observer_gain(self):
        """
        计算Luenberger观测器的增益矩阵
        使用极点配置方法，基于扩展系统的可观测性
        """
        # 检查扩展系统的可观测性
        n = self.A_sq.shape[0]  # 扩展系统的状态维度
        O = self.C_sq.copy()
        temp = self.C_sq.copy()

        # 构建可观测性矩阵
        for i in range(1, n):
            temp = temp @ self.A_sq
            O = np.vstack((O, temp))

        # 验证可观测性
        rank = np.linalg.matrix_rank(O)
        if rank < n:
            print(f"Warning: System not fully observable. Rank: {rank}/{n}")

        # 设置期望极点
        # 将极点分为两组：
        # 1. 状态估计相关的极点 - 较快的动态响应
        # 2. 扰动估计相关的极点 - 较慢的动态响应
        poles_state = np.linspace(-4, -2, 12)  # 状态估计极点
        poles_dist = np.linspace(-1, -0.5, 3)  # 扰动估计极点
        desired_poles = np.concatenate([poles_state, poles_dist])

        # 使用极点配置方法计算观测器增益
        L = control.place(self.A_sq.T, self.C_sq.T, desired_poles).T

        return L

    def update(self, u, y):
        exd_state = np.block([[self.x_hat],
                              [self.d_hat]])
        exd_state = self.A_sq @ exd_state + \
                    self.B_sq @ u.reshape(-1, 1) + \
                    self.L @ (y.reshape(-1, 1) - self.C_sq @ exd_state) - self.gravity
        self.x_hat = exd_state[:12]
        self.d_hat = exd_state[12:]

# %% start simulation
quad_model = Quadrotor()
quad_model.reset()
real_trajectory = {'x': [], 'y': [], 'z': []}
des_trajectory = {'x': [], 'y': [], 'z': []}

accu_error_pos = np.zeros((3,))
total_time = 0
square_ang_vel = np.zeros((4,))

simu_freq = 100  # Hz
ctrl_freq = 50
traj = Trajectory("circle")

quad_controller = Linear_MPC(traj, ctrl_freq, use_obsv=False)

simu_time = 40  # sec
cur_time = 0
dt = 1 / simu_freq
num_iter = int(simu_time * simu_freq)
start = time.time()
for i in range(num_iter):
    des_state = traj.get_des_state(cur_time)
    if i % (int(simu_freq / ctrl_freq)) == 0:
        control_input, error_pos = quad_controller.control(cur_time, quad_model.state)
    cmd_rotor_speeds = control_input["cmd_motor_speeds"]
    obs, _, _, _ = quad_model.step(cmd_rotor_speeds)
    cur_time += dt

    accu_error_pos += error_pos * dt
    square_ang_vel += cmd_rotor_speeds ** 2 * dt

    real_trajectory['x'].append(obs['x'][0])
    real_trajectory['y'].append(obs['x'][1])
    real_trajectory['z'].append(obs['x'][2])

    des_trajectory['x'].append(des_state['x'][0][0])
    des_trajectory['y'].append(des_state['x'][1][0])
    des_trajectory['z'].append(des_state['x'][2][0])

'''Print three required criterions'''
print("Tracking performance: ", np.sum(accu_error_pos ** 2))
print("Sum of square of angular velocities: ", np.sum(square_ang_vel))
print("Total time: ", time.time() - start)

'''Visualization'''
visualizer = Visualizer(simu_time, simu_freq, ctrl_freq, real_trajectory, des_trajectory)
visualizer.plot_tracking_performance()
try:
    if quad_controller.use_obsv == True:
        visualizer.plot_obsv_x(quad_controller.x_real, quad_controller.x_obsv)
except:
    pass
visualizer.plot_3d_trajectory()